# play_animation.py
import numpy as np
from rpi5_ws2812.ws2812 import Color, WS2812SpiDriver, Strip
import time
import os
from config import config

# 初始化LED
driver = WS2812SpiDriver(spi_bus=0, spi_device=0, led_count=512)
strip = Strip(driver)

# 帧参数
LED_COUNT = 512
FRAME_SIZE = LED_COUNT * 3

def load_animation_file(filename):
    if not os.path.exists(filename):
        raise FileNotFoundError(f"动画文件不存在: {filename}")

    file_size = os.path.getsize(filename)
    frames = file_size // FRAME_SIZE

    if file_size % FRAME_SIZE != 0:
        print(f"警告: 文件大小不为 {FRAME_SIZE} 的倍数")

    if config.DEBUG_MODE:
        print(f"加载动画: {filename}, 共 {frames} 帧, 大小: {file_size / 1024:.1f}KB")

    return frames, filename

def play_animation(filename, frames):
    try:
        with open(filename, 'rb') as f:
            frame_idx = 0.0
            last_time = time.time()
            frame_buffer = bytearray(FRAME_SIZE)

            while True:
                current_frame = int(frame_idx) % frames
                f.seek(current_frame * FRAME_SIZE)

                if f.readinto(frame_buffer) != FRAME_SIZE:
                    print("读取帧失败")
                    break

                frame_data = np.frombuffer(frame_buffer, dtype=np.uint8).reshape(LED_COUNT, 3)
                frame_data = (frame_data * config.BRIGHTNESS).astype(np.uint8)

                for i in range(LED_COUNT):
                    r, g, b = frame_data[i]
                    strip.set_pixel_color(i, Color(r, g, b))
                strip.show()

                # 控制帧率
                elapsed = time.time() - last_time
                time.sleep(max(0, 1.0 / config.TARGET_FPS - elapsed))

                frame_idx += config.ROTATION_SPEED
                last_time = time.time()

    except KeyboardInterrupt:
        strip.clear()
        strip.show()
        print("播放终止")
    except Exception as e:
        print(f"播放错误: {e}")
        strip.clear()
        strip.show()

def main():
    if config.ANIMATION_MODE == "projection":
        filename = f"./vox-bin/{config.VOX_FILE}_projection.bin"
    elif config.ANIMATION_MODE == "y_slice":
        filename = f"./vox-bin/{config.VOX_FILE}_yslice.bin"
    elif config.ANIMATION_MODE == "z_y0_slice":
        filename = f"./vox-bin/{config.VOX_FILE}_z_y0slice.bin"
    else:
        raise ValueError(f"无效动画模式: {config.ANIMATION_MODE}")

    frames, filename = load_animation_file(filename)
    print(f"开始播放: {filename}")
    play_animation(filename, frames)

if __name__ == "__main__":
    main()
